'''
Created on Sep 25, 2009

@author: mkiyer
'''
import operator
import sys
import math
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import numpy as np
import glob
import collections
from optparse import OptionParser


def get_hill_func(Vmax, Km, n):
    '''return a specific hill function with parameters Vmax, Km, n'''
    return lambda s: (float(Vmax) * (s**n))/((Km**n) + (s**n))

def get_inv_hill_func(Vmax, Km, n):
    '''return a specific inverse hill function with parameters Vmax, Km, n'''
    return lambda s: float(Vmax)/((Km**n) + (s**n))
    

class UTRCoverage(object):
    
    def calc_score(self):
        # set a lower bound on coverage value to avoid division by
        # zero errors when calculating the relative expression
        zero_func = lambda x: 0.1 if x < 0.1 else x
        e1 = zero_func(self.e_proximal)
        e2 = zero_func(self.e_distal)
        c1 = zero_func(self.c_proximal)
        c2 = zero_func(self.c_distal)
        # compute relative expression
        self.relexpr = math.log((e1/e2)/(c1/c2))

        # now return a score that assesses the quality of
        # the expression.  to calculate the score we factor
        # in both relative expression and coverage
        
        # require some minimal RPKM within the coding region of the 
        # 3'UTR
        if self.e_cds_cov < 2.0 or self.c_cds_cov < 2.0:
            return 0.0
        return self.relexpr
    
    def calc_score_old(self):        
        # make the variables short to reduce amount of typing,
        # and also change all "zeros" to ones so the log function
        # does not break
        zero_func = lambda x: 1.0 if x < 1.0 else x
        et = zero_func(self.e_total)
        e1 = zero_func(self.e_before)
        e2 = zero_func(self.e_after)
        ct = zero_func(self.c_total)
        c1 = zero_func(self.c_before)
        c2 = zero_func(self.c_after)

        # store the relative expression change
        self.relexpr = math.log((e1/e2)/(c1/c2))

        # impose some hard limits
        if self.e_before < 1.0 or self.c_before < 1.0 or \
            self.e_after < 0.5 or self.c_after < 0.5:
            return 0.0

        # score relative expression
        relexpr_score = math.log((e1/e2)/(c1/c2))

        # transform coverage to squash outliers onto range (0,1)
#        cov_transform = get_hill_func(1, 200, 4)
#        te1 = cov_transform(e1)
#        te2 = cov_transform(e2)
#        tc1 = cov_transform(c1)
#        tc2 = cov_transform(c2)
        #relexpr = math.log((te1/te2)/(tc1/tc2))
        #relexpr_func = get_hill_func(4, 1.0, 2)
        #relexpr_score = relexpr_func(relexpr)        

        # score the coverage
#        cov_func = get_hill_func(0.25, 200, 2)
#        inv_cov_func = get_inv_hill_func(0.25, 200, 2)
#        c_cov_score = cov_func(c1) + cov_func(c2)
#        e_cov_score = cov_func(e1) + inv_cov_func(e2)

        score = relexpr_score
        #+ c_cov_score + e_cov_score
        return score

    def __init__(self, thisfields):
        '''
        symbol
        accession
        chrom
        strand
        cds_exon_start
        cds_exon_end
        exper_cds_cov
        control_cds_cov
        utr_start
        utr_end
        polya_site
        exper_proximal
        exper_distal
        control_proximal
        control_distal
        '''
        self.symbol = thisfields[0]
        self.accession = thisfields[1]
        self.chrom = thisfields[2]
        self.strand = thisfields[3]
        self.cds_start = int(thisfields[4])
        self.cds_end = int(thisfields[5])
        self.e_cds_cov = float(thisfields[6])
        self.c_cds_cov = float(thisfields[7])
        self.start = int(thisfields[8])
        self.end = int(thisfields[9])
        self.polya = int(thisfields[10])
        self.e_proximal = float(thisfields[11])
        self.e_distal = float(thisfields[12])
        self.c_proximal = float(thisfields[13])
        self.c_distal = float(thisfields[14])
        self.mirna_sites = thisfields[15].split(',')
        # calculate score
        self.score = self.calc_score()
        return

    def __str__(self):
        thisutr = [self.symbol,
                   self.accession,
                   self.chrom,
                   self.strand,
                   str(self.cds_start),
                   str(self.cds_end),
                   '%.3f' % self.e_cds_cov,
                   '%.3f' % self.c_cds_cov,
                   str(self.start),
                   str(self.end),
                   str(self.polya),
                   '%s:%d-%d' % (self.chrom, self.start, self.end),
                   '%.3f' % self.e_proximal,
                   '%.3f' % self.e_distal,
                   '%.3f' % self.c_proximal,
                   '%.3f' % self.c_distal,
                   '%s' % ','.join(self.mirna_sites)]
        return '\t'.join(thisutr)

def read_utr_cov_file(fhd):
    '''
    symbol
    accession
    chrom
    strand
    cds_exon_start
    cds_exon_end
    exper_cds_cov
    control_cds_cov
    utr_start
    utr_end
    polya_site
    exper_proximal
    exper_distal
    control_proximal
    control_distal
    '''
    for thisline in fhd:        
        thisline = thisline.strip()
        if thisline is None:
            continue
        if thisline.startswith('#'):
            continue
        thisfields = thisline.split('\t')
        yield UTRCoverage(thisfields)


def histogram(utrcovs):    
    
    data = []
    for utrcov in utrcovs:
        if utrcov.relexpr != 0:
            data.append(utrcov.relexpr)

    data = np.array(data)
    print 'average', np.average(data)

    fig = plt.figure()
    ax = fig.add_subplot(111)
    nbins=250

    # the histogram of the data
    n, bins, patches = ax.hist(data, bins=nbins, facecolor='green', alpha=0.5, log=True, normed=1)

    plt.axvline(linewidth=4, color='r')

    plt.xlabel('Relative expression')
    plt.ylabel('log(Frequency)')
    
    #n, bins, patches = ax.hist(x2, bins=bins, facecolor='red', alpha=0.25)    
    #l, = plt.plot(bins, mlab.normpdf(bins, 0.0, 1.0), 'r--', label='fit', linewidth=3)
    #legend([l, patches[0]], ['fit', 'hist'])

    ax.grid(True)
    plt.show()
    return


def get_output_header():
    return ('\t'.join(['#rank',
                       'relative_3utr_expr',
                       'chrom',
                       'start',
                       'end',
                       'name',
                       'score',
                       'strand',
                       'polya',
                       'c_total',
                       'c_before_polya',
                       'c_after_polya',
                       'e_total',
                       'e_before_polya',
                       'e_after_polya',
                       'mirna_sites']))

if __name__ == '__main__':
    utrcovs = list(read_utr_cov_file(open(sys.argv[1])))    
    utrcovs = sorted(utrcovs, key=operator.attrgetter('score'), reverse=False)
    report_first_n = 20

    #histogram(utrcovs)
    print get_output_header()
#    print '#rank\tscore\tlog2_rel_expr\tname'
    for i, x in enumerate(utrcovs):
        print '%d\t%f\t%s' % (i, x.relexpr, x)
        #print '%d\t%.2f\t%.2f\t%s' % (i, x.score, x.relexpr, x.name)
        #if i == report_first_n:
        #    break

    sys.exit(0)
#    optionparser = OptionParser("usage: %prog [options] <files>")    
#    optionparser.add_option("-a", dest="all",
#                            help="process and combine results from a list of files")
#    optionparser.add_option("-t", dest="threshold", type="float",
#                            help="relative expr threshold")
#    (options, args) = optionparser.parse_args()

    outfhd = sys.stdout
    utrcovs = collections.defaultdict(lambda: [])    
    outfhd.write('name')
    for f in sys.argv[1:]:
        outfhd.write('\t%s' % f)
        for thisutr in read_utr_cov_file(open(f)):
            key = (thisutr.chrom, thisutr.start, thisutr.end, thisutr.polya, thisutr.name)
            utrcovs[key].append(thisutr)
    outfhd.write('\n')

    for k,utrlist in utrcovs.iteritems():
        found_zero = False
        for thisutr in utrlist:
            if thisutr.score == 0.0:
                found_zero = True
                break
        if found_zero:
            continue

        outfhd.write('%s' % k[-1])
        for thisutr in utrlist:
            outfhd.write('\t%.2f' % thisutr.score)
        outfhd.write('\n')

#        sorted_utrcovs = sorted(utrcovs.values(), 
#                                key=operator.attrgetter('score'), 
#                                reverse=False)
#        for i, x in enumerate(sorted_utrcovs):
#            print '%d\t%f\t%s' % (i, x.score, x)
